import os

import torch
from torch import nn
from cnn.model import CNN
import matplotlib.pyplot as plt
from cnn.func import train, test, loss_value_plot
from cnn.preprocess import preprocess
import sys
from pathlib import Path
from PyQt6.QtWidgets import (
    QWidget,
    QLineEdit,
    QTextEdit,
    QGridLayout,
    QApplication,
    QPushButton,
    QProgressBar,
    QFileDialog,
    QSpinBox,
    QLabel,
    QDoubleSpinBox,
)
from PyQt6.QtCore import QThread, pyqtSignal, QSize
from PyQt6.QtGui import QTextCursor, QPixmap


class Worker(QThread):
    process = pyqtSignal(str)
    process2 = pyqtSignal(int)
    file_path = ""
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    epochs = 10
    lr = 1e-3
    momentum = 0.9

    def terminate(self):
        super().terminate()

    def start(self):
        super().start()
        self.process.emit("start!")

    def run(self):
        (
            device,
            train_dataloader,
            test_dataloader,
            train_data,
            test_data,
            X_dimension,
            y_dimension,
        ) = preprocess(self.file_path)
        CNN_model = CNN(y_dimension, Path(self.file_path).name)
        CNN_model.to(device=self.device)
        epochs = self.epochs
        lr = self.lr
        momentum = self.momentum
        optimizer = torch.optim.SGD(CNN_model.parameters(), lr=lr, momentum=momentum)
        loss_fn = nn.CrossEntropyLoss()
        if os.path.exists("CNN_model.pth") and os.path.exists("CNN_loss.png"):
            try:
                CNN_model.load_state_dict(
                    torch.load("CNN_model.pth", weights_only=True)
                )
            except Exception:
                Path("CNN_model.pth").unlink()
                losses, iter = train(
                    CNN_model,
                    optimizer,
                    loss_fn,
                    epochs,
                    train_dataloader,
                    device,
                    X_dimension,
                    train_data,
                    self.process,
                    self.process2,
                )
            torch.save(CNN_model.state_dict(), "CNN_model.pth")
            loss_value_plot(losses, iter)
            plt.savefig("CNN_loss.png")
        else:
            losses, iter = train(
                CNN_model,
                optimizer,
                loss_fn,
                epochs,
                train_dataloader,
                device,
                X_dimension,
                train_data,
                self.process,
                self.process2,
            )
            torch.save(CNN_model.state_dict(), "CNN_model.pth")

            loss_value_plot(losses, iter)
            plt.savefig("CNN_loss.png")
        self.process2.emit(888)
        test(
            CNN_model,
            loss_fn,
            test_dataloader,
            device,
            X_dimension,
            test_data,
            self.process,
            self.process2,
        )


class Example(QWidget):

    def __init__(self):
        super().__init__()

        self.initUI()

    def initUI(self):

        self.select_file = QPushButton("data file")
        self.file_path = QLineEdit(
            str(
                Path(__file__)
                .parent.joinpath("data")
                .joinpath("kddcup_10precent.zst")
                .absolute()
            )
        )
        self.file_path.setReadOnly(True)
        self.train_btn = QPushButton("start")
        self.stop_btn = QPushButton("stop")

        self.progress = QProgressBar(self)

        self.epochs_label = QLabel("epochs:")
        self.lr_label = QLabel("lr:")
        self.momentum_label = QLabel("momentum:")
        self.pixmap = QPixmap()

        self.pixmap_label = QLabel()
        self.pixmap_label.setFixedSize(QSize(640, 480))
        self.epochs_input = QSpinBox()
        self.epochs_input.setMinimum(1)
        self.epochs_input.setValue(10)

        self.lr_input = QDoubleSpinBox()
        self.lr_input.setSingleStep(0.001)
        self.lr_input.setDecimals(4)
        self.lr_input.setValue(1e-3)
        self.momentum_input = QDoubleSpinBox()
        self.momentum_input.setSingleStep(0.1)
        self.momentum_input.setValue(0.9)

        self.text_line = QTextEdit()
        self.text_line.setReadOnly(True)
        self.select_file.clicked.connect(self.Select_file)
        self.train_btn.clicked.connect(self.Train)
        self.epochs_input.valueChanged.connect(self.update_epochs)
        self.lr_input.valueChanged.connect(self.update_lr)
        self.momentum_input.valueChanged.connect(self.update_momentum)
        self.stop_btn.clicked.connect(self.Stop)

        grid = QGridLayout()
        grid.setSpacing(10)
        subgrid = QGridLayout()

        grid.addWidget(self.select_file, 1, 0)
        grid.addWidget(self.file_path, 1, 1)

        grid.addWidget(self.epochs_label, 2, 0)
        grid.addWidget(self.epochs_input, 2, 1)
        grid.addWidget(self.lr_label, 3, 0)
        grid.addWidget(self.lr_input, 3, 1)
        grid.addWidget(self.momentum_label, 4, 0)
        grid.addWidget(self.momentum_input, 4, 1)

        subgrid.addWidget(self.train_btn, 1, 0)
        subgrid.addWidget(self.stop_btn, 1, 1)
        grid.addLayout(subgrid, 5, 0)
        grid.addWidget(self.progress, 5, 1)
        grid.addWidget(self.text_line, 6, 0)
        grid.addWidget(self.pixmap_label, 6, 1)

        self.setLayout(grid)
        self.setGeometry(300, 300, 900, 600)
        self.setWindowTitle("CNN for kddcup99")
        self.worker = Worker()
        self.worker.process.connect(self.update_line_edit)
        self.worker.process2.connect(self.update_progress)

    def Select_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "select model file", "All Files(*)"
        )
        if file_path:
            self.file_path.setText(file_path)
            self.worker.file_path = file_path

    def Train(self):
        self.worker.file_path = self.file_path.text()
        self.worker.start()

    def Stop(self):
        self.worker.terminate()

    def update_line_edit(self, text):
        if text == "terminate!":
            self.text_line.clear()
            self.text_line.setText(text)
            return
        if text == "start!":
            self.text_line.clear()
        self.text_line.append(text)
        text = self.text_line.toPlainText()
        lines = text.split("\n")
        if len(lines) > 30:
            lines = lines[1:]  # 移除最上面的一行
            self.text_line.setPlainText("\n".join(lines))
            cursor = self.text_line.textCursor()
            cursor.movePosition(QTextCursor.MoveOperation.End)
            self.text_line.setTextCursor(cursor)

    def update_progress(self, progress_num):
        if progress_num >= 0 and progress_num <= 100:
            self.progress.setValue(progress_num)
        elif progress_num == 999:
            print(999)
            self.text_line.clear()
        elif progress_num == 888:
            print(888)
            self.pixmap.load("CNN_loss.png")
            self.pixmap_label.setPixmap(self.pixmap)
            self.pixmap_label.show()

        else:
            pass

    def update_epochs(self, epochs):
        self.worker.epochs = epochs

    def update_lr(self, lr):
        self.worker.lr = lr

    def update_momentum(self, momentum):
        self.worker.momentum = momentum


def main():

    app = QApplication(sys.argv)
    ex = Example()
    ex.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    pass
